# 股票信息管理类
import pandas as pd


class Bar:
    """一个用来获取股票数据并对股票数据进行简单处理的类"""

    def __init__(self, code, indicatorlist):
        """
        :param code:股票代码
        :type code:str
        :param indicatorlist:策略所需要的因子
        :type indicatorlist: list

        """
        self.code = code
        self.indicatorlist = indicatorlist
        self.stock_data = None

    def get_After_day(self, all_data, DateTime):  # 获取往后最近交易日期
        DateList = sorted(list(all_data.keys()))
        if DateTime > DateList[-1]:
            print("%s 日期信息无法找到匹配日期" % DateTime)
        elif DateTime not in DateList:
            for aDateTime in DateList:
                if aDateTime > DateTime:
                    DateTime = aDateTime
                    break
            return DateTime
        else:
            return DateTime

    def get_Previous_day(self, all_data, DateTime):   # 获取往前最近交易日期
        DateList = sorted(list(all_data.keys()), reverse=True)
        if DateTime < DateList[-1]:
            print("%s 日期信息无法找到匹配日期" % DateTime)
        elif DateTime not in DateList:
            for aDateTime in DateList:
                if aDateTime < DateTime:
                    DateTime = aDateTime
                    break
            return DateTime
        else:
            return DateTime

    def selectStockInformation(self, all_data, startDateTime, endDateTime):

        startDateTime = self.get_After_day(all_data, startDateTime)  # 更新起始日期
        endDateTime = self.get_Previous_day(all_data, endDateTime)  # 更新结束日期

        stocklist = []  # 将所需股票信息存储成list
        date_index = [] # 用来存储交易日期
        if startDateTime and endDateTime:
            for key, value in all_data.items():  # 对原始数据进行循环，key是时间，value是所有股票信息（Dataframe）
                if startDateTime <= key <= endDateTime and self.code in value.index:
                    astockdata = value.loc[self.code]         # 索引股票所在行
                    informationlist = astockdata[self.indicatorlist].values.tolist()  # 索引股票所需因子数据
                    stocklist.append(informationlist)
                    date_index.append(key)

        if len(stocklist) != 0:
            # 创建单只股票信息的DataFrame
            self.stock_data = pd.DataFrame(stocklist, index=date_index, columns=self.indicatorlist)
        else:
            print("信息输入有误：请检查时间范围和股票代码")

    def dropNaN(self, axis=0):
        # 去除带有NaN的数据，默认按行
        if self.stock_data.empty:
            print("该股票还没有信息，请先获取股票信息")
        else:
            self.stock_data = self.stock_data.dropna(axis=axis, how="any")

    def getMaxorMin(self, indicator, day_count, type=0):

        """获取股票回测期间的某个因子的前N天的最大值或者最小值序列

        :param indicator: 选择的某个因子
        :param day_count: 前N天
        :param type: 0,Max; 1,Min

        """
        if self.stock_data.empty:
            print("该股票还没有信息，请先获取股票信息")
        else:
            indicator_data = self.stock_data[indicator]   # 获取所选因子的列数据，series类
            list = []
            for i in range(0, len(indicator_data)):
                if i < day_count:
                    list.append(float("nan"))    # 前面N天数据不够，设为NaN
                else:
                    if type == 0:   # 获取最大值
                        list.append(indicator_data[i-day_count:i].max())
                    if type == 1:   # 获取最小值
                        list.append(indicator_data[i-day_count:i].min())
            series = pd.Series(list, index=self.stock_data.index.tolist())
            return series
